#!/usr/bin/env python3
"""
Babeltrace MCP Server

面向 babeltrace 命令行工具的 MCP 服务器，提供以下工具：

1. babeltrace_help - 显示 babeltrace 帮助信息
   参数：无
   返回：{"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}

2. babeltrace_list_formats - 列出并验证可用格式
   参数：无
   返回：{"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
"""

import subprocess
import json
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("Babeltrace MCP Server")

@mcp.tool()
def babeltrace_help() -> dict:
    """
    显示 babeltrace 命令的帮助信息
    
    返回：
        dict: 包含执行结果的字典，结构为：
            {
                "success": bool,      # 命令是否成功执行
                "command": str,       # 执行的命令
                "exit_code": int,     # 退出码
                "stdout": str,        # 标准输出
                "stderr": str         # 标准错误
            }
    """
    cmd = ["babeltrace", "--help"]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def babeltrace_list_formats() -> dict:
    """
    列出 babeltrace 可用的格式并验证特定格式是否存在
    
    返回：
        dict: 包含执行结果的字典，结构为：
            {
                "success": bool,      # 命令是否成功执行
                "command": str,       # 执行的命令
                "exit_code": int,     # 退出码
                "stdout": str,        # 标准输出
                "stderr": str         # 标准错误
            }
    """
    cmd = ["babeltrace", "-l"]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        # 检查输出中是否包含预期的格式列表
        expected_formats = "ctf, dummy, lttng-live, ctf-metadata, text"
        format_check = expected_formats in result.stdout
        
        return {
            "success": result.returncode == 0 and format_check,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()